from typing import Callable
from fractions import Fraction

import numpy as np
from numpy.typing import NDArray
from numpy import pi, sin, cos, sqrt, trace as Tr, real as Re, abs as Abs
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.lines as lines
from matplotlib import ticker

__all__ = ['Square', 'composite_pulses', 'plot_robust_landscape', 'compare_robust_landscape', '_angle2str'
	, 'repr_composite_pulse']

I, sx, sy, sz = np.identity(2, dtype=complex), np.array([[0., 1], [1, 0]], dtype=complex), np.array(
	[[0, -1.j], [1.j, 0.]]), np.array([[1, 0], [0, -1.]], dtype=complex)

def Square(theta: float, phi: float, w: float, g: float, f: float) -> NDArray:
	T: float = theta / w
	w1: float = (1.0 + g) * w
	factor = sqrt(w1 ** 2 + f ** 2)
	if np.isclose(factor, 0.0): return I.copy()
	half_theta: float = factor * T / 2.0
	return cos(half_theta) * I - 1.j * sin(half_theta) / factor * (sx * cos(phi) * w1 + sy * sin(phi) * w1 + f * sz)

def Gauss(theta: float, phi: float, w: float, g: float, f: float, N: int = 250) -> NDArray:
	sigma = theta / (np.sqrt(2 * pi) * w)
	T = 9.0 * sigma
	ts, dt = np.linspace(0., T, N, retstep=True)
	ple_uc = (1.0 + g) * w * np.exp(-(ts - T / 2) ** 2 / (2 * sigma ** 2))
	U = I.copy()
	for w1 in ple_uc:
		factor = sqrt(w1 ** 2 + f ** 2)
		if np.isclose(factor, 0.): continue
		half_theta: float = factor * dt / 2.0
		prop = cos(half_theta) * I - 1.j * sin(half_theta) / factor * (sx * cos(phi) * w1 + sy * sin(phi) * w1 + f * sz)
		U = prop @ U
	return U

def composite_pulses(thetas: NDArray, phis: NDArray) -> Callable:
	def pulse(Rx: Callable, theta, phi, w: float, g: float, f: float):
		wgf: (float, float, float) = (w, g, f)
		U: NDArray = I.copy()
		for theta, phi in zip(thetas, phis): U = Rx(theta, phi, *wgf) @ U
		return U

	return pulse

def get_infide_meas(infide_type: str = 'corr2', **kwargs):
	match infide_type.lower():
		case 'norm2':
			infide_func = lambda UF, UT: \
				Re(Tr((UF - UT).T.conj() @ (UF - UT)))
		case 'norm2 upto global phase':
			infide_func = lambda UF, UT: Re(Tr(UF.T.conj() @ UF) + Tr(UT.T.conj() @ UT) - 2 * Abs(Tr(UF.T.conj() @ UT)))
		case 'corr2':
			P = kwargs.get('P', np.identity(2, dtype=complex))
			infide_func = lambda UF, UT: \
				np.clip(1.0 - Abs(Tr(P @ UF.T.conj() @ UT) / Tr(P)) ** 2, 1e-12, 1.0)
		case 'corr2 upto z phase':
			def infide_func(UF, UT):
				U = UF.T.conj() @ UT
				z_phase = np.angle(U[0, 0]) - np.angle(U[1, 1])
				Uz = cos(z_phase / 2) * I - 1.j * sin(z_phase / 2) * sz
				return np.clip(1.0 - Abs(Tr(UF.T.conj() @ UT @ Uz) / 2) ** 2, 1e-12, 1.0)
		case _:
			raise NotImplementedError
	return infide_func

def get_infide_mat(theta: float, phi: float, w: float, gs: NDArray, fs: NDArray, ideal_gate: Callable,
                   pulse_gate: Callable, infide_type: str = 'corr2', **kwargs) -> NDArray:
	infide_mat: NDArray = np.zeros((len(fs), len(gs)))
	infide_func = get_infide_meas(infide_type, **kwargs)
	ideal_Rx: NDArray = ideal_gate(theta, phi, w, 0., 0.)
	for i, g in enumerate(gs):
		for j, f in enumerate(fs):
			real_Rx: NDArray = pulse_gate(theta, phi, w, g, f)
			infide_mat[j, i] = infide_func(ideal_Rx, real_Rx)
	return infide_mat

def _plot_robust_landscape(g_mat: NDArray, f_mat: NDArray, infide_ARR: NDArray, fig=None, ax=None, cmap=None, norm=None,
                           cb=False):
	if fig is None: fig, ax = plt.subplots(1, 1, figsize=(2.2, 2.1), layout='constrained')

	if norm is None: norm = mpl.colors.LogNorm(1e-6, 1e-1)
	if cmap is None: cmap = mpl.colormaps['viridis'].reversed().resampled(100)
	lev_exp = np.power(10, np.linspace(-6, -1, 60))

	cts = ax.contourf(g_mat, f_mat, infide_ARR, levels=lev_exp, norm=norm, cmap=cmap, extend='both')
	ax.contour(g_mat, f_mat, infide_ARR, levels=[1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1],
	           colors=['red', 'red', 'k', 'grey', 'white', 'orange'], linewidths=1.0,
	           linestyles=[':', '-', '--', '-', '--', ':'])
	if cb:
		fig.colorbar(cts, ax=ax, orientation='vertical',
		             fraction=.05, pad=0.05, aspect=25, extend='both',
		             ticks=[1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6],
		             format=mpl.ticker.FixedFormatter(
			             ['$10^{-1}$', '$10^{-2}$', '$10^{-3}$', '$10^{-4}$', '$10^{-5}$', '$10^{-6}$']))

def plot_robust_landscape(theta: float, phi: float, w: float, gs: NDArray, fs: NDArray, ideal_gate: Callable,
                          pulse_gate: Callable, infide_type='corr2', fig=None, ax=None, cmap=None, norm=None,
                          cb=False, **kwargs):
	if fig is None: fig, ax = plt.subplots(1, 1, figsize=(2.2, 2.1), layout='constrained')
	g_ARR, f_ARR = np.meshgrid(gs, fs)
	infide_ARR = get_infide_mat(theta, phi, w, gs, fs, ideal_gate, pulse_gate, infide_type=infide_type, **kwargs)
	_plot_robust_landscape(g_ARR, f_ARR, infide_ARR, fig=fig, ax=ax, cmap=cmap, norm=norm, cb=cb)

def _angle2str(angle: float, precision: int = 4) -> str:
	angle = round(angle / pi, precision)
	num, den = Fraction(str(angle)).as_integer_ratio()
	if den == 1:
		return r'0' if num == 0 else r'\pi' if num == 1 else r'-\pi' if num == -1 else f'%s\pi' % num
	elif den <= 24:
		return r'\frac{%s}{%s}\pi' % (num, den) if num > 0 else r'-\frac{%s}{%s}\pi' % (-num, den)
	else:
		return f'{angle:.{precision}f}\pi'

def compare_robust_landscape(theta: float, phi: float, w: float, gs: NDArray, fs: NDArray, ideal_gate: Callable,
                             pulse_gates: dict, infide_type='corr2', fig=None, axs=None, legend=False, **kwargs):
	len_pulses: int = len(pulse_gates)
	if fig is None:
		fig = plt.figure(figsize=(2.5 * len_pulses, 3.0), layout='constrained')
	if axs is None:
		axs = fig.subplots(1, len_pulses)
	for i, (pulse_name, pulse_gate) in enumerate(pulse_gates.items()):
		ax = axs[i] if len_pulses > 1 else axs
		plot_robust_landscape(theta, phi, w, gs, fs, ideal_gate, pulse_gate, fig=fig, ax=ax,
		                      infide_type=infide_type, cb=True if i == len_pulses - 1 else False, **kwargs)
		ax.set_xlabel(r'$\beta_\Omega$')
		ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, pos: f'{x / w:.1f}'))
		ax.set_title(pulse_name)
	ax = axs[0] if len_pulses > 1 else axs
	ax.set_ylabel(r'$\beta_z/\Omega_m$')
	for i in range(1, len(pulse_gates)):
		ax = axs[i] if len_pulses > 1 else axs
		ax.yaxis.set_ticklabels([])
	if legend:
		labels, colors, linestyles = ['$10^{-1}$', '$10^{-2}$', '$10^{-3}$', '$10^{-4}$', '$10^{-5}$'], \
			['orange', 'w', 'grey', 'k', 'red'], [':', '--', '-', '--', '-']
		fig.legend(handles=[lines.Line2D([0], [0], label=label, c=c, ls=ls) for label, c, ls in
		                    zip(labels, colors, linestyles)], ncols=1, loc='lower left', facecolor='silver',
		           shadow=True)

def repr_composite_pulse(thetas, phis, new_line=4, precision=4):
	repr_str = ''
	for i, (theta, phi) in enumerate(zip(thetas, phis)):
		theta_as_str, phi_as_str = _angle2str(theta, precision), _angle2str(phi, precision)
		pre_str = r"\rightarrow" if i != 0 else ''
		repr_str += pre_str + r"(%s)_{%s}" % (theta_as_str, phi_as_str)
		if i != 0 and i % new_line == 0: repr_str += '\n'
	return repr_str
